from typing import Union
from simplejwt import util
from simplejwt.jwt import default_alg, _hash, Jwt
import datetime
import json


class DateEncoder(json.JSONEncoder):
    """
    重写该方法是为了解决json.dumps()在序列化datetime.datetime对象时的问题
    :param obj: 需要序列化的对象
    :return: 序列化后的字符串
    """
    def default(self, obj):
        if isinstance(obj, datetime.datetime):
            return obj.strftime("%Y-%m-%d %H:%M:%S")
        else:
            return json.JSONEncoder.default(self, obj)


def encode(
    secret: Union[str, bytes],  # 密钥，可以是字符串或字节类型
    payload: dict = None,  # 载荷，可选的字典参数
    alg: str = default_alg,  # 算法，默认值为default_alg
    header: dict = None,  # 头部，可选的字典参数
) -> str:
    secret = util.to_bytes(secret)  # 将密钥转换为字节类型

    payload = payload or {}  # 如果载荷未提供，则设置为空字典
    header = header or {}  # 如果头部未提供，则设置为空字典

    header_json = util.to_bytes(json.dumps(header))  # 将头部字典转换为JSON字符串，并将其转换为字节类型
    header_b64 = util.b64_encode(header_json)  # 对头部JSON字符串进行Base64编码
    payload_json = util.to_bytes(
        json.dumps(payload, cls=DateEncoder)
    )  # 将载荷字典转换为JSON字符串，并将其转换为字节类型
    payload_b64 = util.b64_encode(payload_json)  # 对载荷JSON字符串进行Base64编码

    pre_signature = util.join(header_b64, payload_b64)  # 将头部和载荷的Base64编码连接起来作为前签名
    signature = _hash(secret, pre_signature, alg)  # 使用密钥、前签名和算法计算签名
    signature_b64 = util.b64_encode(signature)  # 对签名进行Base64编码

    token = util.join(pre_signature, signature_b64)  # 将前签名和签名Base64编码连接起来生成令牌
    return util.from_bytes(token)  # 将令牌从字节类型转换为字符串类型并返回


class Jwt(Jwt):
    def encode(self) -> str:
        payload = {}
        payload.update(self.registered_claims)
        payload.update(self.payload)
        return encode(self.secret, payload, self.alg, self.header)